package org.example.classloader.ASM;

import jdk.internal.org.objectweb.asm.ClassReader;
import jdk.internal.org.objectweb.asm.ClassWriter;
import org.example.classloader.FileClassLoader;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Set;

public class TransactionClassLoader extends FileClassLoader {

    public TransactionClassLoader(String rootDir) {
        super(rootDir);
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {

        byte[] classData = super.getClassData(name);
        Class<?> target;
        if (classData == null) {
            throw new ClassNotFoundException();
        }

        Set<String> tranMehodSet = null;
        InputStream inputStream;

        // 读取之前的class 字节码
        ClassReader classReader = null;
        try {
            inputStream = new ByteArrayInputStream(classData);
            tranMehodSet = getTransSet(inputStream);
        } catch (IOException e) {
            e.printStackTrace();
        }

        if (null == tranMehodSet) {
            target = super.defineClass(name, classData, 0, classData.length);
            return target;
        }
        // 进行字节码的解析
        try {
            inputStream = new ByteArrayInputStream(classData);
            classReader = new ClassReader(inputStream);
        } catch (IOException e) {
            e.printStackTrace();
        }
        ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS);
        classReader.accept(new ModifyClassVisitor(classWriter, tranMehodSet), ClassReader.SKIP_DEBUG);
        // 取得解析后的字节码
        byte[] transactionClassBits = classWriter.toByteArray();
        target = defineClass(name, transactionClassBits, 0, transactionClassBits.length);
        System.out.println("src class=" + target.getName());
        return target;
    }

    private static Set<String> getTransSet(InputStream inputStream) throws IOException {

        ClassReader classReader = new ClassReader(inputStream);
        FilterClassVisitor filterClassVisitor = new FilterClassVisitor();
        classReader.accept(filterClassVisitor, ClassReader.SKIP_DEBUG);
        Set<String> tranMehodSet = filterClassVisitor.getTranMehods();
        return tranMehodSet;
    }
}